import akka.actor.{Actor, ActorSystem, PoisonPill, Props}
import scala.concurrent.duration.DurationInt

case class Register(key: String)
case class Beat(key: String)
case object Ack

class Master extends Actor{
  val DURATION_THRESHOLD = 5 // seconds

  val workers = scala.collection.mutable.HashMap[String, Long]()

  val thread = new Thread {
    override def run(): Unit = {
      println("start master check loop....")

      while (true) {
        Thread.sleep(1000)
        checkTimestamp()
      }
    }
  }
  thread.start

  override def receive: Receive = {
    case Register(key) =>
      if (workers.contains(key)) {
        println(s"worker[$key] register conflict!")
      } else {
        updateTimestamp(key)
        println(s"worker[$key] register successful.")
        sender() ! Ack
      }
    case Beat(key) =>
      if (workers.contains(key)) { // 响应 worker 心跳，更新时间戳
        updateTimestamp(key)
        sender() ! Ack
      } else { // 关闭超时的worker
        sender() ! PoisonPill
      }

  }

  private def updateTimestamp(key: String) = {
    this.synchronized {
      workers(key) = System.currentTimeMillis()
    }
  }

  private def checkTimestamp(): Unit = {
    this.synchronized {
      val now = System.currentTimeMillis()

      println("检查心跳时间戳：")
      for ((worker, ts) <- workers.toSet) {
        val duration = (now - ts) / 1000
        println(s"$worker => $duration 秒")
        if (duration > DURATION_THRESHOLD) {
          println(s"$worker 已经超时")
          workers.remove(worker)
        }
      }
    }
  }

}

object Master{
  def main(args: Array[String]): Unit = {
    val system = ActorSystem("App")
    val master = system.actorOf(Props[Master], "master")

    import system.dispatcher

    system.scheduler.scheduleOnce(500 millis) {
      for(i <- 1 to 5) {
        val name = s"worker-$i"
        val worker = system.actorOf(Props(classOf[Slaver], master, name), name)
        worker ! Register
      }
    }
  }
}
